import requests
import time
import os
from typing import Generator
import pprint

import jwt

from data_types import TextMsg, ImageMsg, TextMsgList, MsgList, CharacterMeta


# 智谱开放平台API key，参考 https://open.bigmodel.cn/usercenter/apikeys
API_KEY: str = os.getenv("API_KEY", "")



class ApiKeyNotSet(ValueError):
    pass


def verify_api_key_not_empty():
    if not API_KEY:
        raise ApiKeyNotSet


def generate_token(apikey: str, exp_seconds: int) -> str:
    # reference: https://open.bigmodel.cn/dev/api#nosdk
    try:
        id, secret = apikey.split(".")
    except Exception as e:
        raise Exception("invalid apikey", e)
 
    payload = {
        "api_key": id,
        "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
        "timestamp": int(round(time.time() * 1000)),
    }
 
    return jwt.encode(
        payload,
        secret,
        algorithm="HS256",
        headers={"alg": "HS256", "sign_type": "SIGN"},
    )




def get_characterglm_response_via_sdk(messages: TextMsgList) -> Generator[str, None, None]:
    """ 通过新版sdk调用characterglm """
    # 与get_characterglm_response等价
    # Reference: https://open.bigmodel.cn/dev/api#characterglm
    # 需要安装旧版sdk，zhipuai==1.0.7
    from zhipuai import ZhipuAI
    verify_api_key_not_empty()
    client = ZhipuAI (api_key=API_KEY)
    # pprint.pprint(messages)
    
    response = client.chat.completions.create(
        model="charglm-4",
        messages = messages
    )

    # pprint.pprint(response.choices[0].message.content)
    return response # .choices[0].message.content
    # for event in response.events():
    #     if event.event == 'add':
    #         yield event.data


def get_chatglm_response_via_sdk(messages: TextMsgList) -> Generator[str, None, None]:
    """ 通过sdk调用chatglm """
    # reference: https://open.bigmodel.cn/dev/api#glm-3-turbo  `GLM-3-Turbo`相关内容
    # 需要安装新版zhipuai
    from zhipuai import ZhipuAI
    verify_api_key_not_empty()
    client = ZhipuAI(api_key=API_KEY) # 请填写您自己的APIKey
    response = client.chat.completions.create(
        model="GLM-4-AirX",  # 填写需要调用的模型名称
        messages=messages,
        stream=False,
    )
    # for chunk in response:
    #     yield chunk.choices[0].delta.content
    return response.choices[0].message.content


def generate_role_appearance(role_profile: str) -> Generator[str, None, None]:
    """ 用chatglm生成角色的外貌描写 """
    
    instruction = f"""
请从下列文本中，抽取人物的外貌描写。若文本中不包含外貌描写，请你推测人物的性别、年龄，并生成一段外貌描写。要求：
1. 只生成外貌描写，不要生成任何多余的内容。
2. 外貌描写不能包含敏感词，人物形象需得体。
3. 尽量用短语描写，而不是完整的句子。
4. 不要超过50字

文本：
{role_profile}
"""
    return get_chatglm_response_via_sdk(
        messages=[
            {
                "role": "user",
                "content": instruction.strip()
            }
        ]
    )

def generate_roles_from_description(description: str) -> Generator[str, None, None]:
    """ 用chatglm生成角色的人设描述 """
    
    instruction = f"""
请从下列文本之后的描述中，抽取人物的人设描述，文本中可能包含多个出场人物，需要为每个有名字的人物各自生成一份人设描述。若文本中不包含人设描述，或者出场的人物没有人设描述，请你推测人物的性别、年龄，并生成一段人设描述。
要求：
1. 只生成人设描述，不要生成任何多余的内容，每个人物的人设描述使用一段文本，不要拆分成多段。
2. 人设描述不能包含敏感词，人物形象需得体。
3. 尽量用短语描写，而不是完整的句子。
4. 不要超过50字。
5. 人物和人设描述之间用中文冒号分隔。

文本内容如下：
{description}
"""
    print(instruction)
    return get_chatglm_response_via_sdk(
        messages=[
            {
                "role": "user",
                "content": instruction.strip()
            }
        ]
    )

def generate_chat_scene_prompt(messages: TextMsgList, meta: CharacterMeta) -> Generator[str, None, None]:
    """ 调用chatglm生成cogview的prompt，描写对话场景 """
    instruction = f"""
阅读下面的角色人设与对话，生成一段文字描写场景。

{meta['bot_name']}的人设：
{meta['bot_info']}
    """.strip()
    
    if meta["user_info"]:
        instruction += f"""

{meta["user_name"]}的人设：
{meta["user_info"]}
""".rstrip()

    if messages:
        instruction += "\n\n对话：" + '\n'.join((meta['bot_name'] if msg['role'] == "assistant" else meta['user_name']) + '：' + msg['content'].strip() for msg in messages)
    
    instruction += """
    
要求如下：
1. 只生成场景描写，不要生成任何多余的内容
2. 描写不能包含敏感词，人物形象需得体
3. 尽量用短语描写，而不是完整的句子
4. 不要超过50字
""".rstrip()
    print(instruction)
    
    return get_chatglm_response_via_sdk(
        messages=[
            {
                "role": "user",
                "content": instruction.strip()
            }
        ]
    )


def generate_cogview_image(prompt: str) -> str:
    """ 调用cogview生成图片，返回url """
    # reference: https://open.bigmodel.cn/dev/api#cogview
    from zhipuai import ZhipuAI
    print('API_KEY:',API_KEY)
    client = ZhipuAI(api_key=API_KEY) # 请填写您自己的APIKey
    
    response = client.images.generations(
        model="cogview-3-flash", #填写需要调用的模型名称
        prompt=prompt
    )
    return response.data[0].url

def combine_roles(role1: str, role2: str) -> str:
    """ 用chatglm生成角色的人设描述 """
    
    instruction = f"""
请将以下两组人设描述中相同的人物的条目合并成一段描述，合并时保留所有人设信息，去除重复内容，确保描述简洁且不超过50字。合并后每个人物的描述之间使用‘|’分隔开。

示例：
组1：
角色1：喜欢吃糖果。
角色2：喜欢吃糖果，喜欢看电影。
组2：
角色2：喜欢看电影。
合并后：
角色1：喜欢吃糖果|角色2：喜欢吃糖果，喜欢看电影。

请按照示例合并以下两组人设描述：
组1：
{role1}
组2：
{role2}
"""
    print(instruction)
    return get_chatglm_response_via_sdk(
        messages=[
            {
                "role": "user",
                "content": instruction.strip()
            }
        ]
    )

def generate_role_play(roles: str, messages: list, topic_description: str) -> str:
    """ 用charglm生成角色对话 """
    if not messages:
        if topic_description:
            instruction = [{"role":"system", "content":f"""你是{roles}，请根据对话记录，生成回答。"""},
                        {"role":"user", "content":topic_description}]
        else:
            instruction = [{"role":"system", "content":f"""你是{roles}，请生成随机话题。"""},
                        {"role":"user", "content":"随机对话"}]
        # pprint.pprint(instruction)
        return get_characterglm_response_via_sdk(
                messages= instruction
                )
    else:
        instruction = [{"role":"system", "content":f"""你是{roles}，请根据对话记录，生成回答。"""}]
        # pprint.pprint(instruction)
        return get_characterglm_response_via_sdk(
                messages= instruction + messages
                )
    
    